"""
返回召回的结果
"""
from dnn.recall.sentence_vectors import Sentence2Vector

import config


class Recall:
    def __init__(self, method="BM25"):
        assert method.lower() in ["bm25", "tfidf", "fasttext"], "无效的方法"
        self.method = method.lower()
        self.sentence = Sentence2Vector(method=self.method)
        self.vectorizer, self.features_vector, self.lines_cut, self.search_index = \
            self.sentence.build_vector()

    def predict(self, sentence):
        # sentence: {"cut": str, "entity": list}
        # tfidf.transform的参数必须是一个列表
        search_vector = self.vectorizer.transform([sentence["cut"]])
        # print(self.features_vector.shape)
        # print(search_vector.shape)
        cp_search_list = self.search_index.search(search_vector, k=config.recall_topk,
                                                  k_clusters=config.recall_clusters,
                                                  return_distance=True)
        # print(cp_search_list)
        # """
        # [[('0.0', '什么是Python？'), ('0.0', '什么是python'), ('0.0', 'python是做什么的？'),
        #  ('0.3300653230435001', 'Python做什么项目？')]]
        # """
        # print(type(cp_search_list))  # <class 'list'>
        # print(cp_search_list[0])
        # """
        # [('0.0', '什么是Python？'), ('0.0', '什么是python'), ('0.0', 'python是做什么的？'),
        # ('0.3300653230435001', 'Python做什么项目？')]
        # """

        # 过滤主体
        filter_results = []
        for result in cp_search_list[0]:
            # result是cp_search_list中的元组
            # print(result)
            # distance = result[0]
            key = result[1]
            entities = self.sentence.qa_dict[key]["entity"]
            if len(set(entities) & set(sentence["entity"])) > 0:
                filter_results.append(result)

        # 返回最终结果
        if len(filter_results) < 1:
            return [i[1] for i in cp_search_list[0]]
        else:
            return [i[1] for i in filter_results]
